#/usr/bin/env python3
# -*- coding: UTF-8 -*-
import numpy as np
import pdb
from PIL import Image
import numbers
file="~/data/quant_data_cv/n01440764/ILSVRC2012_val_00009346.JPEG"

model_conf = {
      "isBGR": True,
      "resize": (256, 256),
      "crop": 224,
      "mergeProcess": True,
      "mean": [104.0, 117.0, 123.0],
      "std": [1.0, 1.0, 1.0],
  }

def resize(img, size, interpolation=Image.BILINEAR):
    if isinstance(size, int):
        w, h = img.size
        if (w <= h and w == size) or (h <= w and h == size):
            return img
        if w < h:
            ow = size
            oh = int(size * h / w)
            return img.resize((ow, oh), interpolation)
        else:
            oh = size
            ow = int(size * w / h)
            return img.resize((ow, oh), interpolation)
    else:
        return img.resize(size, interpolation)

def center_crop(img, output_size):
    if isinstance(output_size, numbers.Number):
        output_size = (int(output_size), int(output_size))
    w, h = img.size
    th, tw = output_size
    i = int(round((h - th) / 2.))
    j = int(round((w - tw) / 2.))
    img = img.crop((j, i, j+tw, i+th))
    return img

def mean_std_process(data):
    mean = np.array(model_conf["mean"])
    mean = mean[:, np.newaxis, np.newaxis]
    std = np.array(model_conf["std"])
    std = std[:, np.newaxis, np.newaxis]
    data = (data-mean)/std
    return data


if __name__ == '__main__':    
    with open(file, 'rb') as f:
        data = Image.open(f)
        data = data.convert("RGB")
        data = resize(data, model_conf['resize'])
        data = center_crop(data, model_conf['crop'])
        data = np.array(data)
        data = data.transpose((2, 0, 1))
        if(model_conf['isBGR']):
            data = data[::-1, :, :]
        if not model_conf['mergeProcess']:
            data = mean_std_process(data)
        dtype = data.dtype
        data.flatten().tofile(f"first_frame_in_data_{dtype}.bin")
        print(f"Success to save {file} -> first_frame_in_data_{dtype}.bin.")

